{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from transformers import AutoTokenizer, AutoConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ckpt = \"bert-base-uncased\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_ckpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([[ 2051, 10029,  2066,  2019,  8612]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}\n"
     ]
    }
   ],
   "source": [
    "text = \"time flies like an arrow\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\", add_special_tokens=False)\n",
    "print(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BertConfig {\n",
      "  \"_name_or_path\": \"bert-base-uncased\",\n",
      "  \"architectures\": [\n",
      "    \"BertForMaskedLM\"\n",
      "  ],\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"classifier_dropout\": null,\n",
      "  \"gradient_checkpointing\": false,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"model_type\": \"bert\",\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"pad_token_id\": 0,\n",
      "  \"position_embedding_type\": \"absolute\",\n",
      "  \"transformers_version\": \"4.40.2\",\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n",
      "Embedding(30522, 768)\n"
     ]
    }
   ],
   "source": [
    "config = AutoConfig.from_pretrained(model_ckpt)\n",
    "print(config)\n",
    "token_emb = nn.Embedding(config.vocab_size, config.hidden_size)\n",
    "print(token_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 2051, 10029,  2066,  2019,  8612]])\n",
      "torch.Size([1, 5, 768])\n",
      "tensor([[[-0.5399,  0.8817, -0.0582,  ...,  0.1344,  1.2879,  1.3271],\n",
      "         [-0.4227,  0.7234,  0.4978,  ...,  0.1828,  0.6424, -0.4907],\n",
      "         [ 1.4265,  1.8380, -0.1233,  ..., -0.0595,  1.7091, -0.4492],\n",
      "         [-0.3932, -0.8089,  2.3442,  ..., -0.7789,  1.8642,  1.2017],\n",
      "         [-0.1489,  0.7013,  0.5352,  ...,  0.3394, -0.5013,  0.2779]]],\n",
      "       grad_fn=<EmbeddingBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(inputs.input_ids)\n",
    "inputs_emb = token_emb(inputs.input_ids)\n",
    "print(inputs_emb.shape)\n",
    "print(inputs_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from math import sqrt\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "768\n",
      "torch.Size([1, 5, 5])\n",
      "tensor([[[ 3.0308e+01, -9.1329e-01,  6.7655e-01,  1.1622e+00,  4.1983e-01],\n",
      "         [-9.1329e-01,  2.4962e+01, -6.1393e-01,  1.2064e+00, -7.4860e-01],\n",
      "         [ 6.7655e-01, -6.1393e-01,  2.7884e+01,  2.9213e-02,  2.8198e-01],\n",
      "         [ 1.1622e+00,  1.2064e+00,  2.9213e-02,  2.9108e+01, -1.0842e-01],\n",
      "         [ 4.1983e-01, -7.4860e-01,  2.8198e-01, -1.0842e-01,  2.6514e+01]]],\n",
      "       grad_fn=<DivBackward0>)\n"
     ]
    }
   ],
   "source": [
    "Q = K = V = inputs_emb\n",
    "dim_k = K.size(-1)\n",
    "print(dim_k)\n",
    "scores = torch.bmm(Q, K.transpose(1, 2)) / sqrt(dim_k)\n",
    "print(scores.size())\n",
    "print(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1., 1., 1.]], grad_fn=<SumBackward1>)\n",
      "tensor([[[1.0000e+00, 2.7585e-14, 1.3525e-13, 2.1982e-13, 1.0463e-13],\n",
      "         [5.7892e-12, 1.0000e+00, 7.8096e-12, 4.8216e-11, 6.8257e-12],\n",
      "         [1.5272e-12, 4.2019e-13, 1.0000e+00, 7.9939e-13, 1.0293e-12],\n",
      "         [7.2985e-13, 7.6281e-13, 2.3506e-13, 1.0000e+00, 2.0483e-13],\n",
      "         [4.6480e-12, 1.4448e-12, 4.0494e-12, 2.7406e-12, 1.0000e+00]]],\n",
      "       grad_fn=<SoftmaxBackward0>)\n"
     ]
    }
   ],
   "source": [
    "weights = F.softmax(scores, dim=-1)\n",
    "print(weights.sum(dim=-1))\n",
    "print(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n",
      "tensor([[[-0.5399,  0.8817, -0.0582,  ...,  0.1344,  1.2879,  1.3271],\n",
      "         [-0.4227,  0.7234,  0.4978,  ...,  0.1828,  0.6424, -0.4907],\n",
      "         [ 1.4265,  1.8380, -0.1233,  ..., -0.0595,  1.7091, -0.4492],\n",
      "         [-0.3932, -0.8089,  2.3442,  ..., -0.7789,  1.8642,  1.2017],\n",
      "         [-0.1489,  0.7013,  0.5352,  ...,  0.3394, -0.5013,  0.2779]]],\n",
      "       grad_fn=<BmmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "atten_outputs = torch.bmm(weights, V)\n",
    "print(atten_outputs.shape)\n",
    "print(atten_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scaled_dot_product_attention(query, key, value, query_mask=None, key_mask=None, mask=None):\n",
    "    dim_k = query.size(-1)\n",
    "    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)\n",
    "    if query_mask is not None and key_mask is not None:\n",
    "        mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, -float(\"inf\"))\n",
    "    weights = F.softmax(scores, dim=-1)\n",
    "    return torch.bmm(weights, value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionHead(nn.Module):\n",
    "    def __init__(self, embed_dim, head_dim):\n",
    "        super().__init__()\n",
    "        self.q = nn.Linear(embed_dim, head_dim)\n",
    "        self.k = nn.Linear(embed_dim, head_dim)\n",
    "        self.v = nn.Linear(embed_dim, head_dim)\n",
    "    \n",
    "    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):\n",
    "        atten_outputs_inter = scaled_dot_product_attention(\n",
    "            self.q(query), self.k(key), self.v(value), query_mask, key_mask, mask\n",
    "        )\n",
    "        return atten_outputs_inter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BertConfig {\n",
    "#   \"_name_or_path\": \"bert-base-uncased\",\n",
    "#   \"architectures\": [\n",
    "#     \"BertForMaskedLM\"\n",
    "#   ],\n",
    "#   \"attention_probs_dropout_prob\": 0.1,\n",
    "#   \"classifier_dropout\": null,\n",
    "#   \"gradient_checkpointing\": false,\n",
    "#   \"hidden_act\": \"gelu\",\n",
    "#   \"hidden_dropout_prob\": 0.1,\n",
    "#   \"hidden_size\": 768,\n",
    "#   \"initializer_range\": 0.02,\n",
    "#   \"intermediate_size\": 3072,\n",
    "#   \"layer_norm_eps\": 1e-12,\n",
    "#   \"max_position_embeddings\": 512,\n",
    "#   \"model_type\": \"bert\",\n",
    "#   \"num_attention_heads\": 12,\n",
    "#   \"num_hidden_layers\": 12,\n",
    "#   \"pad_token_id\": 0,\n",
    "#   \"position_embedding_type\": \"absolute\",\n",
    "#   \"transformers_version\": \"4.40.2\",\n",
    "#   \"type_vocab_size\": 2,\n",
    "#   \"use_cache\": true,\n",
    "#   \"vocab_size\": 30522\n",
    "# }\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        embed_dim=config.hidden_size\n",
    "        num_heads = config.num_attention_heads\n",
    "        head_dim = embed_dim // num_heads\n",
    "        self.heads = nn.ModuleList(\n",
    "            [AttentionHead(embed_dim=embed_dim, head_dim=head_dim) for _ in range(num_heads)]\n",
    "        )\n",
    "        self.output_linear = nn.Linear(embed_dim, embed_dim)\n",
    "    \n",
    "    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):\n",
    "        x = torch.cat([h(query, key, value, query_mask, key_mask, mask) for h in self.heads], dim=-1)\n",
    "        x = self.output_linear(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n",
      "tensor([[[-0.0790, -0.0361, -0.2813,  ..., -0.0701, -0.0428,  0.1751],\n",
      "         [-0.0936,  0.0693, -0.2840,  ..., -0.0421, -0.0745,  0.1116],\n",
      "         [-0.0850,  0.0354, -0.2509,  ..., -0.0103, -0.0990,  0.1704],\n",
      "         [-0.0543,  0.0237, -0.2928,  ..., -0.0425, -0.1176,  0.1305],\n",
      "         [-0.0520,  0.0034, -0.2697,  ..., -0.0374, -0.0995,  0.2186]]],\n",
      "       grad_fn=<ViewBackward0>)\n"
     ]
    }
   ],
   "source": [
    "multihead_attn = MultiHeadAttention(config)\n",
    "query = key = value = inputs_emb\n",
    "atten_output = multihead_attn(query, key, value)\n",
    "print(atten_output.shape)\n",
    "print(atten_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.linear_1 = nn.Linear(config.hidden_size, config.intermediate_size)\n",
    "        self.linear_2 = nn.Linear(config.intermediate_size, config.hidden_size)\n",
    "        self.gelu = nn.GELU()\n",
    "        self.droput = nn.Dropout(config.hidden_dropout_prob)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.linear_1(x)\n",
    "        x = self.gelu(x)\n",
    "        x = self.linear_2(x)\n",
    "        x = self.droput(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n",
      "tensor([[[ 0.0134,  0.0055,  0.0226,  ...,  0.0769,  0.0000, -0.0444],\n",
      "         [-0.0028, -0.0270,  0.0376,  ...,  0.0000,  0.0249, -0.0000],\n",
      "         [ 0.0120,  0.0029,  0.0256,  ...,  0.0733,  0.0365, -0.0000],\n",
      "         [ 0.0013, -0.0149,  0.0364,  ...,  0.0841,  0.0301, -0.0410],\n",
      "         [ 0.0000, -0.0111,  0.0346,  ...,  0.0824,  0.0384, -0.0473]]],\n",
      "       grad_fn=<MulBackward0>)\n"
     ]
    }
   ],
   "source": [
    "feed_forward = FeedForward(config)\n",
    "ff_outputs = feed_forward(atten_output)\n",
    "print(ff_outputs.size())\n",
    "print(ff_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoderLayer(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.layer_norm_1 = nn.LayerNorm(config.hidden_size)\n",
    "        self.layer_norm_2 = nn.LayerNorm(config.hidden_size)\n",
    "        self.attention = MultiHeadAttention(config)\n",
    "        self.feed_forward = FeedForward(config)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        # Apply layer normalization and then copy input into query, key, value\n",
    "        hidden_state = self.layer_norm_1(x)\n",
    "        # Apply attention with a skip connection\n",
    "        x = x + self.attention(hidden_state, hidden_state, hidden_state, mask=mask)\n",
    "        # Apply feed-forward layer with a skip connection\n",
    "        x = x + self.feed_forward(self.layer_norm_2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n",
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "source": [
    "encoder_layer = TransformerEncoderLayer(config)\n",
    "print(inputs_emb.shape)\n",
    "print(encoder_layer(inputs_emb).size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "source": [
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.token_embeddings = nn.Embedding(config.vocab_size,\n",
    "                                             config.hidden_size)\n",
    "        self.position_embeddings = nn.Embedding(config.max_position_embeddings,\n",
    "                                                config.hidden_size)\n",
    "        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, input_ids):\n",
    "        # Create position IDs for input sequence\n",
    "        seq_length = input_ids.size(1)\n",
    "        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0)\n",
    "        # Create token and position embeddings\n",
    "        token_embeddings = self.token_embeddings(input_ids)\n",
    "        position_embeddings = self.position_embeddings(position_ids)\n",
    "        # Combine token and position embeddings\n",
    "        embeddings = token_embeddings + position_embeddings\n",
    "        embeddings = self.layer_norm(embeddings)\n",
    "        embeddings = self.dropout(embeddings)\n",
    "        return embeddings\n",
    "\n",
    "embedding_layer = Embeddings(config)\n",
    "print(embedding_layer(inputs.input_ids).size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.embeddings = Embeddings(config)\n",
    "        self.layers = nn.ModuleList([TransformerEncoderLayer(config)\n",
    "                                     for _ in range(config.num_hidden_layers)])\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        x = self.embeddings(x)\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask=mask)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "source": [
    "encoder = TransformerEncoder(config)\n",
    "print(encoder(inputs.input_ids).size())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
